import torch.utils.data as data
from PIL import Image
import PIL
import os
import os.path
import pickle
import random
import numpy as np
import pandas as pd


class TextDataset(data.Dataset):
    def __init__(self, data_dir, split='train'):
        self.imsize = 64
        self.data = []
        self.data_dir = data_dir
        split_dir = os.path.join(data_dir, split)

        self.filenames = self.load_filenames(split_dir)
        self.embeddings = self.load_embedding(split_dir)

    def get_img(self, img_path):
        img = Image.open(img_path).convert('RGB')
        load_size = int(self.imsize * 76 / 64)
        img = img.resize((load_size, load_size), PIL.Image.BILINEAR)  # 图像resize+插值
        return img

    def load_embedding(self, data_dir):
        embedding_path = os.path.join(data_dir, 'embeddings.pkl')
        with open(embedding_path, 'rb') as f:
            embeddings = pickle.load(f)
            embeddings = np.array(embeddings)
            print('embeddings: ', embeddings.shape)
        return embeddings

    def load_filenames(self, data_dir):
        filepath = os.path.join(data_dir, 'filenames.pkl')
        with open(filepath, 'rb') as f:
            filenames = pickle.load(f)
        print('Load filenames from: %s (%d)' % (filepath, len(filenames)))
        return filenames

    def __getitem__(self, index):
        key = self.filenames[index]
        data_dir = self.data_dir
        embedding = self.embeddings[index, :]
        img_name = '%s/image/%s' % (data_dir, key)
        print(img_name)
        img = self.get_img(img_name)

        return img, embedding

    def __len__(self):
        return len(self.filenames)


data_dir = 'D:/NLP_txt2img'
dataset = TextDataset(data_dir, 'train')
img_, embedding_ = dataset.__getitem__(2)

print(type(dataset.embeddings), np.shape(dataset.embeddings))
print(type(dataset.filenames), np.shape(dataset.filenames))

print(type(img_), np.shape(img_))
print(img_)
print(type(embedding_), np.shape(embedding_))
print(embedding_)
